from flask import Flask,request,send_file
from flask_cors import CORS
from pic2svg import rings
from draw import drawin_create
import shutil
import time
import os
# import draw
import sd
import host
# from pic2svg import rings
from draw import drawin_create
from PIL import Image, ImageTk
import argparse
import socket

import subprocess 
def tkUI():
    command = ['python', './src/mainUI.py']
    subprocess.call(command)
parser = argparse.ArgumentParser(
                    prog='ProgramName',
                    description='What the program does',
                    epilog='Text at the bottom of help')
parser.add_argument('-i', action='store_true') 
args = parser.parse_args()
if args.i:
    tkUI()

host.link()




app = Flask(__name__)
cors = CORS(app, resources={r"/*": {"origins": "*"}})
pic_name = ''
status_code = 0

last_beat = 0.0
last_client = ''



@app.route('/post', methods=["POST"])
def post_data():
    global status_code
    global pic_name
    p = request.get_json() 
    r = {}
    if p['command']=='prompt': #
        prompt = p['prompt']
        # gen_pic = sd_generate(prompt)
        negative_prompt = "low quality"
        url = "http://127.0.0.1:7860"
        sd.sd_work(prompt, negative_prompt, url)
        gen_pic = "./tmp/pic_output/output.png"
        name = '/sd/'+time.strftime("%Y-%m-%d_%H-%M-%S")+'.png'
        pic_name = './src/webui'+name
        shutil.copy(gen_pic,pic_name)
        r ={
            'msg':'prompt OK',
            'pic_url':'.'+name
        }

    elif p['command']=='pic_trans':
        # trans_pic = pic_trans(pic_name)
        prompt = p['prompt']
        negative_prompt = "low quality"
        url = "http://127.0.0.1:7860"
        print(prompt)
        sd.img2img(prompt, negative_prompt, url, pic_name)
        trans_pic = "./tmp/pic_output/output.png"
        name = '/sd/'+time.strftime("%Y-%m-%d_%H-%M-%S")+'.trans.png'
        pic_name = './src/webui'+name
        shutil.copy(trans_pic,pic_name)
        r ={
            'msg':'pic_trans OK',
            'pic_url':'.'+name
        }
        pass

    elif p['command']=='code_generate': #指令生成
        rings(pic_name, './tmp/temp_file/img_canny.svg',
            blur_size=int(p['blur']),layers_gate=[5000,3000,1200,int(p['gate'])])
        drawin_create('./tmp/temp_file/img_canny.svg',no_te=True)
        name = '/svg_png/'+time.strftime("%Y-%m-%d_%H-%M-%S")+'.svg.png'
        shutil.copy('./tmp/temp_file/img_canny.svg.png','./src/webui'+name)
        r = {
            'msg':'code_gen OK',
            'pic_url':'.'+name
        }
    
    elif p['command']=='run_pause': #启动/暂停/继续
        host.go()
        if status_code == 0:
            status_code = 1
        else: status_code = 3 - status_code
        r['msg']='run_pause OK'

    elif p['command']=='stop': #停止
        host.stop()
        status_code = 0
        r['msg']='stop OK'
        pass

    elif p['command']=='status_check':
        r = {'code':status_code}
    
    elif p['command']=='tune':
        direction = p['direction']
        if direction=='up':
            host.origin[2]+=0.5
        elif direction=='down':
            host.origin[2]-=0.5
        r['msg']='tune OK'
    
    elif p['command']=='pic_clean':
        pic_clean()
        pic_name = ''
        r['msg']='clean OK'
    return r

@app.route('/file_upload', methods=["POST"])
def file_get():
    f = request.files['file']
    ext = f.filename.split('.')[1]
    name = time.strftime("%Y-%m-%d_%H-%M-%S")+'.'+ext
    f.save('./src/webui/upload/'+name)
    global pic_name
    pic_name = './src/webui/upload/' + name
    r = {
        'pic_url':'./upload/' + name
    }
    return r

@app.route('/heartbeat', methods=["POST"])
def heart_beat():
    global last_beat, last_client, status_code
    p = request.get_json()
    r = {}
    t = time.time()
    if not host.t.is_alive():
        status_code = 0
    if t - last_beat >= 5:
        last_beat = t
        last_client = p['user_id']
        r['status']='OK'
    else:
        if p['user_id'] == last_client:
            last_beat = t
            r['status']='OK'
        else:
            r['status']='Bad'
    return r

@app.route('/', methods=["GET"])
def index():
    a = send_file('./webui/index.html')
    return a

@app.route('/',defaults={'path':''},methods=['GET'])
@app.route('/<path:path>', methods=["GET"])
def file_server(path):
    a = send_file(f'./webui/{path}')
    return a

def pic_clean():
    path_list = [
        './src/webui/svg_png',
        './src/webui/upload',
        './src/webui/sd'
    ]
    for path in path_list:
        file_list = os.listdir(path)
        for f in file_list:
            os.remove(path+'/'+f)

if __name__ == '__main__':
    pic_clean()
    
    app.run(debug=False, host='0.0.0.0', port=80)